import onnx_graphsurgeon as gs
import onnx
import numpy as np


def make_onnx(qdq=False,
              num_batch=4,
              num_queries=11273,
              num_pts=8,
              num_heads=8,
              num_hidden_dim=32,
              shape=[[128, 240], [64, 120], [32, 60], [16, 30]],
              scale_input_sampling_locations=0.0078125,
              scale_input_attention_weights=0.0078125,
              scale_input_values=[0.01, 0.01, 0.01, 0.01],
              scale_outputs=0.02,
              feature_layout='NHWC'):
    graph = gs.Graph()

    input_sampling_locations = gs.Variable(
        name='input_sampling_locations',
        shape=[num_batch, num_queries, num_heads, len(shape), num_pts, 2],
        dtype=np.float32)
    input_attention_weights = gs.Variable(
        name='input_attention_weights',
        shape=[num_batch, num_queries, num_heads, len(shape), num_pts],
        dtype=np.float32)
    input_values = []
    msda = gs.Node(op='MSDA', attrs={'feature_layout': feature_layout})
    if feature_layout in ['NHWC', 'NCHW32']:
        for si, s in enumerate(shape):
            feat_shape = [num_batch, num_heads*num_hidden_dim] + s
            input_values.append(gs.Variable(name=f"input_feat_{si}",
                                shape=feat_shape, dtype=np.float32))
    elif feature_layout in ['CONCAT_NHWC', 'CONCAT_NCHW32']:
        msda.attrs['spatial_shape'] = sum(shape, [])
        spatial_locations = sum([s[0]*s[1] for s in shape])
        if feature_layout == 'CONCAT_NHWC':
            feat_shape = [num_batch, spatial_locations,
                          num_heads*num_hidden_dim]
        elif feature_layout == 'CONCAT_NCHW32':
            assert num_heads*num_hidden_dim % 32 == 0
            feat_shape = [num_batch, num_heads *
                          num_hidden_dim // 32, spatial_locations, 32]
        input_values.append(gs.Variable(name=f"input_feat_concat",
                            shape=feat_shape, dtype=np.float32))
    else:
        raise NotImplementedError

    output = gs.Variable(name=f"output", dtype=np.float32)

    graph.inputs = [input_sampling_locations,
                    input_attention_weights] + input_values
    graph.outputs = [output]

    if qdq:
        input_sampling_locations_Q = gs.Variable(
            name=f'{input_sampling_locations.name}_Q')
        quant_input_sampling_locations = gs.Node(
            op='QuantizeLinear',
            name=f'quant_{input_sampling_locations.name}',
            inputs=[input_sampling_locations,
                    gs.Constant(name=f'scale_{input_sampling_locations.name}', values=np.array(
                        scale_input_sampling_locations, dtype=np.float32)),
                    gs.Constant(name=f'zp_{input_sampling_locations.name}', values=np.array(0, dtype=np.int8))],
            outputs=[input_sampling_locations_Q])

        input_attention_weights_Q = gs.Variable(
            name=f'{input_attention_weights.name}_Q')
        quant_input_attention_weights = gs.Node(
            op='QuantizeLinear',
            name=f'quant_{input_attention_weights.name}',
            inputs=[input_attention_weights,
                    gs.Constant(name=f'scale_{input_attention_weights.name}', values=np.array(
                        scale_input_attention_weights, dtype=np.float32)),
                    gs.Constant(name=f'zp_{input_attention_weights.name}', values=np.array(0, dtype=np.int8))],
            outputs=[input_attention_weights_Q])

        input_values_Q = []
        quant_input_values = []
        for v, scale_v in zip(input_values, scale_input_values):
            v_Q = gs.Variable(name=f'{v.name}_Q')
            input_values_Q.append(v_Q)
            quant_input_values.append(
                gs.Node(op='QuantizeLinear',
                        name=f'quant_{v.name}',
                        inputs=[v,
                                gs.Constant(name=f'scale_{v.name}', values=np.array(
                                    scale_v, dtype=np.float32)),
                                gs.Constant(name=f'zp_{v.name}', values=np.array(0, dtype=np.int8))],
                        outputs=[v_Q]))

        output_Q = gs.Variable(name=f'{output.name}_Q')
        dequant_output = gs.Node(op='DequantizeLinear',
                                 name=f'dequant_output',
                                 inputs=[output_Q,
                                         gs.Constant(name=f'scale_{output.name}', values=np.array(
                                             scale_outputs, dtype=np.float32)),
                                         gs.Constant(name=f'zp_{output.name}', values=np.array(0, dtype=np.int8))],
                                 outputs=[output])

        msda.name = 'MSDA_qdq'
        msda.inputs = [input_sampling_locations_Q,
                       input_attention_weights_Q] + input_values_Q
        msda.outputs = [output_Q]
        msda.attrs['scale_value'] = [q.inputs[1].values.item()
                                     for q in quant_input_values]
        msda.attrs['scale_offset'] = quant_input_sampling_locations.inputs[1].values.item()
        msda.attrs['scale_weight'] = quant_input_attention_weights.inputs[1].values.item()
        msda.attrs['scale_out'] = dequant_output.inputs[1].values.item()

        graph.nodes = [msda, quant_input_sampling_locations,
                       quant_input_attention_weights, dequant_output] + quant_input_values

    else:
        msda.name = 'MSDA'
        msda.inputs = [input_sampling_locations,
                       input_attention_weights] + input_values
        msda.outputs = [output]
        graph.nodes = [msda]

    graph.cleanup()
    onnx.save(gs.export_onnx(graph),
              f'msda_{"QDQ" if qdq else "noQDQ"}_{feature_layout}.onnx')

    np.random.rand(*input_sampling_locations.shape).astype(
        np.float32).tofile('input_sampling_locations.bin')


if __name__ == '__main__':
    make_onnx(qdq=False, feature_layout='NHWC')
    make_onnx(qdq=False, feature_layout='CONCAT_NHWC')
    make_onnx(qdq=False, feature_layout='CONCAT_NCHW32')
    make_onnx(qdq=True, feature_layout='NCHW32')
